#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
快速端口关闭工具 - 一键关闭8000-8010端口
"""

import subprocess
import psutil
import sys
import os


def get_processes_on_ports(start_port=8000, end_port=8010):
    """获取占用指定端口范围的进程"""
    processes = []
    
    for port in range(start_port, end_port + 1):
        try:
            # 使用netstat命令
            cmd = f'netstat -ano | findstr ":{port} "'
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
            
            if result.stdout.strip():
                lines = result.stdout.strip().split('\n')
                for line in lines:
                    if f':{port}' in line and 'LISTENING' in line:
                        parts = line.split()
                        pid = int(parts[-1])
                        try:
                            process = psutil.Process(pid)
                            processes.append({
                                'pid': pid,
                                'port': port,
                                'name': process.name(),
                                'cmdline': ' '.join(process.cmdline()) if process.cmdline() else ''
                            })
                        except:
                            processes.append({
                                'pid': pid,
                                'port': port,
                                'name': 'Unknown',
                                'cmdline': ''
                            })
        except Exception as e:
            print(f"查询端口 {port} 时出错: {e}")
    
    return processes


def kill_process(pid, force=False):
    """结束进程"""
    try:
        process = psutil.Process(pid)
        if force:
            process.kill()
        else:
            process.terminate()
        
        try:
            process.wait(timeout=3)
        except psutil.TimeoutExpired:
            if not force:
                process.kill()
                process.wait(timeout=2)
        
        return True
    except Exception as e:
        print(f"结束进程 {pid} 失败: {e}")
        return False


def main():
    """主函数"""
    print("快速端口关闭工具")
    print("=" * 40)
    
    # 检查psutil
    try:
        import psutil
    except ImportError:
        print("安装 psutil...")
        subprocess.run([sys.executable, "-m", "pip", "install", "psutil"], check=True)
        import psutil
    
    # 获取占用端口的进程
    processes = get_processes_on_ports(8000, 8010)
    
    if not processes:
        print("端口 8000-8010 范围内没有程序在监听")
        return
    
    print(f"发现 {len(processes)} 个进程占用端口:")
    for proc in processes:
        print(f"  端口 {proc['port']}: PID {proc['pid']} - {proc['name']}")
        if proc['cmdline']:
            print(f"    命令: {proc['cmdline']}")
    
    # 直接关闭所有进程
    print("\n正在结束所有进程...")
    killed = 0
    for proc in processes:
        if kill_process(proc['pid'], force=True):
            killed += 1
            print(f"已结束 PID {proc['pid']} ({proc['name']})")
    
    print(f"\n完成！共结束 {killed} 个进程")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n程序中断")
    except Exception as e:
        print(f"错误: {e}")